import { useState, forwardRef, useRef, useImperativeHandle } from 'react';
import { Button, Flex, message } from 'antd';
import { mockStreamAPI } from './service';
import { Bubble, Sender } from '@ant-design/x';

type AIChatBoxProps = {
  apiKey: string;
  initialMessage: string;
};

type Message = {
  role: 'user' | 'assistant';
  content: string;
};

type AIChatBoxRef = {
  resetFlag: () => void;
};

const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms))
const AIChatBox = forwardRef<AIChatBoxRef, AIChatBoxProps>(({ apiKey, initialMessage }, ref) => {
  const [messages, setMessages] = useState<Message[]>(initialMessage ? [{ role: 'assistant', content: initialMessage }] : [])
  const [isLoading, setIsLoading] = useState<boolean>(false)
  const [input, setInput] = useState<string>('')

  useImperativeHandle(ref, () => ({
    resetFlag: () => {
      setMessages((prev) => {
        if (prev.length === 1) {
          message.info('无需清空')
        } else {
          message.success('清空成功')
        }
        return initialMessage ? [{ role: 'assistant', content: initialMessage }] : []
      })
    }
  }), [initialMessage, setMessages])

  const handleSubmit = async (message: string) => {
    setInput('')
    setIsLoading(true)
    setMessages([...messages, { role: 'user', content: message }, { role: 'assistant', content: '' }])
    try {
      // 模拟网络请求, 展示loading
      await sleep(1000)
      const response = mockStreamAPI(apiKey, message)
      for await (const item of response) {
        // 捕获接口返回的错误信息
        if (item.error) {
          setMessages(prev => {
            const lastMessage = prev[prev.length - 1]
            return [...prev.slice(0, -1), { ...lastMessage, content: lastMessage.content + item.error }]
          })
          break;
        }
        setMessages(prev => {
          const lastMessage = prev[prev.length - 1]
          return [...prev.slice(0, -1), { ...lastMessage, content: lastMessage.content + item.char }]
        })
        await sleep(50)  // 模拟流式返回
      }
      // 捕获服务端或未知错误
    } catch (error) {
      setMessages(prev => {
        const lastMessage = prev[prev.length - 1]
        return [...prev.slice(0, -1), { ...lastMessage, content: lastMessage.content + '模型服务繁忙，请重试' }]
      })
    }
    setIsLoading(false)
  }

  const chatItems = messages.map((message, index) => ({
    role: message.role,
    content: message.content,
    key: message.role + index,
    // 判断助手消息是否正在加载
    loading: message.role === 'assistant' && isLoading && index === messages.length - 1 && message.content === '',
  }))

  return <Flex vertical style={{ minHeight: '50vh', border: '1px solid #ccc', borderRadius: 10, padding: 20 }} justify='space-between'>
    <Bubble.List items={chatItems} style={{ maxHeight: 'calc(50vh - 66px)', overflowY: 'auto' }} roles={{
      user: {
        placement: 'end',
      },
      assistant: {
        placement: 'start',
      },
    }}></Bubble.List>
    {/*注意：针对"题目二"的要求三 这个组件自带防抖功能，所以不需要我们再手动防抖，如果是自己手搓组件，需要自己手动添加防抖。 */}
    <Sender placeholder="问我任何问题" loading={isLoading} value={input} onChange={(value) => setInput(value)} onSubmit={handleSubmit} />
  </Flex>;
});


const Home = () => {

  const ref = useRef<AIChatBoxRef>(null)

  return <div style={{ margin: 20 }}>
    <Button style={{ marginBottom: 10 }} onClick={() => ref.current?.resetFlag()}>清空</Button>
    <AIChatBox ref={ref} apiKey="123" initialMessage="你好 我是AI助手" />
  </div>
}

export default Home